import sys
import logging
import numpy as np
import tensorflow as tf
import h5py


def create_correct_model_architecture():
    model = tf.keras.layers.Sequential([
        tf.keras.layers.Embedding(
            input_dim=256,
            output_dim=128,
            input_length=2048,
            name='embedding'
        ),
        tf.keras.layers.LSTM(
            256,
            return_sequences=False,
            name='lstm'
        ),
        tf.keras.layers.Dense(
            5,
            activation='softmax',
            name='dense'
        )
    ])

    model.build(input_shape=(None, 2048))
    model.compile(
        optimizer='adam',
        loss='sparse_categorical_crossentropy',
        metrics=['accuracy']
    )

    return model

def load_weights_from_h5(model_path):
    weights = {}

    with h5py.File(model_path, 'r') as f:
        mw = f['model_weights']

        weights['embedding'] = np.array(mw['embedding/embedding/embeddings:0'])

        #
        weights['lstm_kernel'] = np.array(mw['lstm/lstm/lstm_cell/kernel:0'])
        weights['lstm_recurrent_kernel'] = np.array(mw['lstm/lstm/lstm_cell/recurrent_kernel:0'])
        weights['lstm_bias'] = np.array(mw['lstm/lstm/lstm_cell/bias:0'])

        #
        weights['dense_kernel'] = np.array(mw['dense/dense/kernel:0'])
        weights['dense_bias'] = np.array(mw['dense/dense/bias:0'])
    
    return weights

def convert_model_final(input, output):
    model = create_correct_model_architecture()
    model.summary()

    weights = load_weights_from_h5(input)
    model.layers[0].set_weights([weights['embedding']])
    logging.info(f"Set embedding weights {weights['embedding'].shape}")

    model.layers[1].set_weights([
        weights['lstm_kernel'],
        weights['lstm_recurrent_kernel'],
        weights['lstm_bias']
    ])

    model.layers[2].set_weights([
        weights['dense_kernel'],
        weights['dense_bias']
    ])

    model.save(output)

    test_model = tf.keras.models.load_model(output)
    test_model.summary()

def main():
    input = 'hapray/optimization_detector/models/aarch64-flag-lstm.h5'
    output = 'hapray/optimization_detector/models/aarch64-flag-lstm-converted.h5'
    convert_model_final(input, output)

if __name__ == '__main__':
    main()